﻿using System;
using System.IO;
using System.Collections.Generic;
using System.Linq;
using UnityEngine;


public class BVHDriver : MonoBehaviour
{
    [Header("Loader settings")]
    [Tooltip("This is the target avatar for which the animation should be loaded. Bone names should be identical to those in the BVH file and unique. All bones should be initialized with zero rotations. This is usually the case for VRM avatars.")]
    public Animator targetAvatar;

    [Serializable]
    public struct BoneMap
    {
        public string name;
        public HumanBodyBones humanoid_bone;
    }

    private BoneMap[] bmapsALL; // the corresponding bones between unity and bvh
    private BoneMap[] bmapsBVH; // the corresponding bones between unity and bvh
    private BoneMap[] bmapsBVHCor; // the corresponding bones between unity and bvh
    private BoneMap[] bmapsNSM; // the corresponding bones between unity and nsm
    private BoneMap[] bmapsNSM_all;
    private BVHParser bp = null;
    private Animator anim;

    private Dictionary<string, Quaternion> bvhT;
    private Dictionary<string, Vector3> bvhOffset;
    private Dictionary<string, string> bvhHireachy;
    private Dictionary<HumanBodyBones, Quaternion> unityT;
    public Dictionary<string, Quaternion> currRotation;

    private int frameIdx;
    public float bvh_scaleRatio = 1.0f;
    public float nsm_scaleRatio = 1.0f;
    private string searchKey = "";
    private float frameRate = 60.0f;
    private float unity_leftleg = 0f;

    public Dictionary<string, Transform> getAllBonesTransform()
    {
        Dictionary<string, Transform> all_bones = new Dictionary<string, Transform>();

        foreach (BoneMap bm in bmapsALL)
        {
            Transform currBone = anim.GetBoneTransform(bm.humanoid_bone);
            all_bones.Add(bm.name, currBone);
        }

        return all_bones;
    }

    private void Start(){
        BonemapReader reader = new BonemapReader();
        string path = Params.bonesMapPath;
        reader.Read(path + "Bonesmap.txt");
        bmapsALL = reader.bonemaps;

        reader.Read(path + "Bonesmap_bvh.txt");
        bmapsBVH = reader.bonemaps;

        reader.Read(path + "Bonesmap_bvh_cor.txt");
        bmapsBVHCor = reader.bonemaps;

        reader.Read(path + "Bonesmap_nsm.txt");
        bmapsNSM = reader.bonemaps;

        reader.Read(path + "Bonesmap_nsm_all.txt");
        bmapsNSM_all = reader.bonemaps;

        // 下面的code，必须从一开始就计算；只有开始时，场景中模型是 T-pose
        anim = targetAvatar.GetComponent<Animator>();
        unityT = new Dictionary<HumanBodyBones, Quaternion>();
        foreach (BoneMap bm in bmapsALL)
        {
            unityT.Add(bm.humanoid_bone,
                anim.GetBoneTransform(bm.humanoid_bone).rotation); // Get maps from unity_bone to its rotations
        }

        unity_leftleg =
            Vector3.Distance(anim.GetBoneTransform(HumanBodyBones.LeftLowerLeg).position,
                anim.GetBoneTransform(HumanBodyBones.LeftUpperLeg).position) +
            Vector3.Distance(anim.GetBoneTransform(HumanBodyBones.LeftLowerLeg).position,
                anim.GetBoneTransform(HumanBodyBones.LeftFoot).position);
    }

    private void search_init()
    {
        // 之后通过 searchkey 进行动作数据库检索，得到 file
        string file = @"Assets\Resources\Bvh\" + searchKey + ".bvh";
        string bvhData = File.ReadAllText(file);
        bp = new BVHParser(bvhData);
        frameRate = 1f / bp.frameTime;
        Application.targetFrameRate = (Int16)frameRate;

        // bvhT = bp.getTFrame();            // Get bones pose of the first frame, string is bone's name, "pos" stands for root
        // bvhT = bp.getKeyFrame(0);         // Get bones pose of the first frame, string is bone's name, "pos" stands for root
        bvhOffset = bp.getOffset(1.0f); // Get bones offset from bvh
        bvhHireachy = bp.getHierachy(); // Get maps from child bone to parent bone

        // unity 与 BVH 骨骼比例
        float bvh_leftleg = 0.0f;
        foreach (BoneMap bm in bmapsALL)
        {
            if (bm.humanoid_bone == HumanBodyBones.LeftLowerLeg || bm.humanoid_bone == HumanBodyBones.LeftFoot)
            {
                bvh_leftleg = bvh_leftleg + (float)Math.Sqrt(bvhOffset[bm.name].sqrMagnitude);
                // unity_leftleg =
                //     Vector3.Distance(anim.GetBoneTransform(HumanBodyBones.LeftLowerLeg).position, anim.GetBoneTransform(HumanBodyBones.LeftUpperLeg).position) + 
                //     Vector3.Distance(anim.GetBoneTransform(HumanBodyBones.LeftLowerLeg).position, anim.GetBoneTransform(HumanBodyBones.LeftFoot).position);
            }
        }

        bvh_scaleRatio = unity_leftleg / bvh_leftleg;

        // unity 与 NSM 骨骼比例
        float nsm_leftleg = 0.9801445f; // 手动得到3D位置计算得到
        nsm_scaleRatio = unity_leftleg / nsm_leftleg;

        Debug.Log("Bone scale ratio[Unity/BVH]: " + bvh_scaleRatio);
        Debug.Log("Bone scale ratio[Unity/NSM]: " + nsm_scaleRatio);
        frameIdx = 4;
    }

    public void searchMotion(string skey)
    {
        if (!searchKey.Equals(skey))
        {
            searchKey = skey;
            if (searchKey == "")
                return;

            // search motion lib by key to get filename
            Debug.Log("Search Style Motion is: " + searchKey);
            search_init();  // initial BVHParser bp            
        }
    }

    private string[] wholeActs = {"punching_79_08", "kick_143_24", "punching_bag", "hook", "pull_heavy", "pushing",
        "hip_hop_dance", "football_124_01", "golf_124_07", "flashlight_77_05", "chopwood_79_01", "lift_79_93"};

    [Header("Pose ScaleRatio")]
    public float scaleRatio = 0.6f;

    // public float adjustRotation = 13f;

    public void updateRotations(Dictionary<string, Quaternion> nsmRots, Vector3 rpos, int mlab)
    {
        // get NSM bone global rotation
        Quaternion rrot = new Quaternion();
        foreach (BoneMap bm in bmapsNSM_all) // Compute global rotations
        {
            Transform currBone = anim.GetBoneTransform(bm.humanoid_bone);
            currBone.rotation = nsmRots[bm.name] * unityT[bm.humanoid_bone];

            if (bm.name == "Head" || bm.name == "Neck")
            {
                currBone.rotation = currBone.rotation * Quaternion.Euler(-15, 0, 0); // 对上半身的角度进行一个调整
            }

            if (bm.name == "Hips")
            {
                rrot = currBone.rotation;
            }
        }

        // 单独处理全身替换动作
        if (mlab == 0 && wholeActs.Contains(searchKey))
        {
            frameIdx = (frameIdx < bp.frames - 1) ? frameIdx + 1 : 4;
            // pull_heavy -> Quaternion.Euler(30, 0, 0)
            // Dictionary<string, Quaternion> currFrame = bp.getKeyFrameByRrot(frameIdx, rrot * Quaternion.Euler(13, 0, 0));

            Quaternion start_bvhRRot = bp.getRootRotation(0);
            Quaternion bvhRRot = bp.getRootRotation(frameIdx);
            Dictionary<string, Quaternion> currFrame = bp.getKeyFrameByRrot(frameIdx, rrot * bvhRRot);
            foreach (BoneMap bm in bmapsALL) // Compute global rotations
            {
                Transform currBone = anim.GetBoneTransform(bm.humanoid_bone);
                currBone.rotation = currFrame[bm.name] * unityT[bm.humanoid_bone]; // unityT must add
            }

            // // set root position
            Vector3 start_bvhRPos = Quaternion.Inverse(start_bvhRRot) * bp.getRootPosition(0);
            Vector3 bvhRPos = Quaternion.Inverse(bvhRRot) * bp.getRootPosition(frameIdx);
            Vector3 actRoot = (bvhRPos - start_bvhRPos) * scaleRatio + rpos;
            anim.GetBoneTransform(HumanBodyBones.Hips).position = actRoot;
            return;
        }

        if (searchKey == "")
        {
            // walk/run/sit, 纠正上身头部位置
            Dictionary<string, Quaternion> currFrame = bp.getKeyFrameByRrot(frameIdx, rrot);
            foreach (BoneMap bm in bmapsBVHCor) // Compute global rotations
            {
                Transform currBone = anim.GetBoneTransform(bm.humanoid_bone);
                currBone.rotation = currFrame[bm.name] * unityT[bm.humanoid_bone]; // unityT must add             
            }
        }
        else
        {
            frameIdx = (frameIdx < bp.frames - 1) ? frameIdx + 1 : 4;
            Dictionary<string, Quaternion> currFrame = bp.getKeyFrameByRrot(frameIdx, rrot * Quaternion.Euler(13, 0, 0));
            foreach (BoneMap bm in bmapsBVH) // Compute global rotations
            {
                Transform currBone = anim.GetBoneTransform(bm.humanoid_bone);
                // currBone.rotation = (currFrame[bm.name] * Quaternion.Inverse(bvhT[bm.name])) * unityT[bm.humanoid_bone]; // unityT must add

                currBone.rotation = currFrame[bm.name] * unityT[bm.humanoid_bone]; // unityT must add
                // currBone.rotation =  currFrame[bm.name]; // unityT must add

                //currBone.rotation = currBone.rotation * Quaternion.EulerAngles(-60, 0, 0);
                // currBone.localRotation = (currFrame[bm.name] * Quaternion.Inverse(bvhT[bm.name])) * unityT[bm.humanoid_bone]; // unityT must add               
            }

            // draw bvh skeleton
            // Dictionary<string, Vector3> bvhPos = new Dictionary<string, Vector3>();
            // foreach (string bname in currFrame.Keys)
            // {
            //    if (bname == "pos")
            //    {
            //     //    bvhPos.Add(bp.root.name, new Vector3(currFrame["pos"].x, currFrame["pos"].y, currFrame["pos"].z));
            //        bvhPos.Add(bp.root.name, new Vector3(rpos.x, currFrame["pos"].y, rpos.z));
            //    }
            //    else
            //    {
            //        if (bvhHireachy.ContainsKey(bname) && bname != bp.root.name)
            //        {
            //            Vector3 curpos = bvhPos[bvhHireachy[bname]] + currFrame[bvhHireachy[bname]] * bvhOffset[bname];
            //            bvhPos.Add(bname, curpos);
            //        }
            //    }
            // }

            // anim.GetBoneTransform(HumanBodyBones.Hips).position = bvhPos[bp.root.name]*scaleRatio;
            // foreach (string bname in bvhHireachy.Keys)
            // {
            //    Color color = new Color(1.0f, 0.0f, 0.0f);
            //    Debug.DrawLine(bvhPos[bname], bvhPos[bvhHireachy[bname]], color);
            // }
        }

        // set root position
        if (mlab == 5) bvh_scaleRatio = 1.05f;
        //Vector3 anim_rpos = new Vector3(rpos.x, rpos.y * bvh_scaleRatio, rpos.z); // 只改变 root height
        Vector3 anim_rpos = new Vector3(rpos.x, rpos.y, rpos.z); // 只改变 root height
        // anim_rpos.y = anim.GetBoneTransform(HumanBodyBones.Hips).position.y;
        anim.GetBoneTransform(HumanBodyBones.Hips).position = anim_rpos;
    }


    // updateOpt
    public Dictionary<string, Matrix4x4> getAllBones()
    {
        Dictionary<string, Matrix4x4> all_bones = new Dictionary<string, Matrix4x4>();

        foreach (BoneMap bm in bmapsALL)
        {
            Transform boneTrans = anim.GetBoneTransform(bm.humanoid_bone);
            Matrix4x4 currBone = Matrix4x4.TRS(boneTrans.position, boneTrans.rotation, boneTrans.localScale);
            all_bones.Add(bm.name, currBone);
        }

        return all_bones;
    }


}

